from sqlalchemy import Column, String, Integer, Float
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

from init.mysql import engine

Base = declarative_base()


class StockSale(Base):
    __tablename__ = 'stock_sale'
    id = Column(Integer, primary_key=True, autoincrement=True)
    code = Column(String(6))
    rdate = Column(String(10))
    ramount = Column(Integer)
    rprice = Column(Float)
    batch = Column(Integer)
    pdate = Column(String(10))

    def __init__(self, code, rdate, ramount, rprice, batch, pdate):
        self.code = code
        self.rdate = rdate
        self.ramount = ramount
        self.rprice = rprice
        self.batch = batch
        self.pdate = pdate

    @staticmethod
    def insert_or_update(code, data_list):
        session = sessionmaker(bind=engine)()
        for data in data_list:
            existing_stock_sale = session.query(StockSale).filter_by(code=code, rdate=data['rdate']).first()
            if existing_stock_sale:
                existing_stock_sale.rdate = data['rdate']
                existing_stock_sale.ramount = data['ramount']
                existing_stock_sale.rprice = data['rprice']
                existing_stock_sale.batch = data['batch']
                existing_stock_sale.pdate = data['pdate']
            else:
                session.add(StockSale(code=code, rdate=data['rdate'], ramount=data['ramount'], rprice=data['rprice'],
                                      batch=data['batch'], pdate=data['pdate']))
        session.commit()


Base.metadata.create_all(engine)
